import { useState, useEffect, useContext } from 'react';
import { Context } from './constant';
import { isEqual, useCustomRef } from './util';

export type ModelItemType = Record<string, any>;

export type Model<T extends keyof ModelItemType> = {
  [key in keyof ModelItemType]: ReturnType<ModelItemType[T]>;
};

export type RetState<
  U extends (...args: any) => any,
  T extends keyof Model<T>,
> = U extends undefined ? Model<T>[T] : ReturnType<U>;

export function useModel<T extends keyof Model<T>, U>(
  key: T,
  selector?: (model: Model<T>[T]) => U,
): RetState<NonNullable<typeof selector>, T> {
  const store = useContext(Context);

  const [getIsMount, setIsMount] = useCustomRef(false);

  const [state, setState] = useState<RetState<NonNullable<typeof selector>, T>>(
    () => (selector ? selector(store.data[key]) : store.data[key]),
  );

  const [getOldState, setOldState] = useCustomRef(state);
  setOldState(state);

  useEffect(() => {
    setIsMount(true);
    const handler = (e: any) => {
      if (getIsMount()) {
        if (selector) {
          const currentState = selector(e);
          if (!isEqual(currentState, getOldState())) {
            setState(currentState);
          }
        } else {
          setState(e);
        }
      }
    };

    store.addHandler(key, handler);

    return () => {
      setIsMount(false);
      store.removeHandler(key, handler);
    };
  }, [key]);

  return state;
}
